package com.gitee.mingjie.okhttp.concurrent;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import okhttp3.Call;
import okhttp3.Callback;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;

/**
 * 实现OkHttp的并发请求
 *
 * @author mingjie
 * @since 2020-11-21
 */
public class ConcurrentHandler {
    private final OkHttpClient okHttpClient;
    private final Map<String, Request> requestMap = new HashMap<>();
    private final Map<String, IOException> exceptionMap = new ConcurrentHashMap<>();
    private final Map<String, Response> responseMap = new ConcurrentHashMap<>();
    private boolean noAwait = Boolean.TRUE;

    public ConcurrentHandler(OkHttpClient okHttpClient) {
        this.okHttpClient = okHttpClient;
    }

    public ConcurrentHandler put(String name, Request request) {
        requestMap.put(name, request);
        return this;
    }

    /**
     * 等待所有请求结束，获取结果前必须调用
     *
     * @return this
     */
    public ConcurrentHandler await() {
        noAwait = Boolean.FALSE;
        CountDownLatch latch = new CountDownLatch(requestMap.size());
        for (Map.Entry<String, Request> requestEntry : requestMap.entrySet()) {
            okHttpClient.newCall(requestEntry.getValue()).enqueue(new Callback() {
                @Override
                public void onFailure(Call call, IOException e) {
                    exceptionMap.put(requestEntry.getKey(), e);
                    latch.countDown();
                }

                @Override
                public void onResponse(Call call, Response response) {
                    responseMap.put(requestEntry.getKey(), response);
                    latch.countDown();
                }
            });
        }
        try {
            latch.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        return this;
    }

    /**
     * 获取所有请求结果，若存在网络失败，则抛出异常，所有可能的异常中的一个
     *
     * @return 请求结果
     */
    public Map<String, Response> completeAll() {
        if (noAwait) {
            throw new RuntimeException("please after calling await");
        }
        if (!exceptionMap.isEmpty()) {
            throw new RuntimeException(exceptionMap.values().iterator().next());
        }

        return responseMap;
    }

    /**
     * 获取部分请求结果（可能存在失败的请求），只有所有网络失败，会抛出异常，所有可能的异常中的一个
     *
     * @return 请求结果
     */
    public Map<String, Response> completeAny() {
        if (noAwait) {
            throw new RuntimeException("please after calling await");
        }
        if (responseMap.isEmpty()) {
            throw new RuntimeException(exceptionMap.values().iterator().next());
        }
        return responseMap;
    }

}
